Hành trình hướng đến các kernel hiệu suất cao bắt đầu bằng việc chuyển đổi từ tập trung vào thao tác lập trình (PyTorch Eager) sang nhận thức được phần cứng lập trình. Triton đóng vai trò là cầu nối then chốt trong hành trình này.
1. Xác định vị trí của hệ thống
Triton là một ngôn ngữ và trình biên dịch cho lập trình song song, được thiết kế để giúp việc viết các kernel tính toán tùy chỉnh hiệu suất cao bằng cú pháp Python trở nên khả thi. Nó nằm ở một vị trí độc đáo giữa hai cực:
- PyTorch Eager: Abstraction cao, dễ sử dụng nhưng kiểm soát hạn chế về việc tận dụng phần cứng.
- CUDA C++: Kiểm soát tối đa, nhưng độ phức tạp cao (quản lý bộ nhớ chung và đồng bộ hóa thủ công).
- Triton: Cú pháp Python với mức khối (chia ô) điều khiển.
2. Mô hình chia ô
Khác với CUDA, hoạt động ở mức luồng, Triton sử dụng mô hình dựa trên khối (chia ô) lập trình. Điều này đặc biệt quan trọng đối với học sâu, nơi dữ liệu (ma trận, bản đồ chú ý) tự nhiên được cấu trúc thành các khối.
3. Sai lầm về hiệu suất
Một hiểu lầm phổ biến là nghĩ rằng Triton chỉ đơn thuần là "PyTorch nhưng nhanh hơn". Thực tế, đó là một mô hình riêng biệt. Các cải thiện hiệu suất đến từ khả năng của nhà phát triển để giải quyết các điểm nghẽn (như "vách tường bộ nhớ") bằng cách hợp nhất các thao tác để giữ dữ liệu trong SRAM trên chip tốc độ cao.